package com.hhh.springai_test.controller;


import com.hhh.springai_test.common.BaseResponse;
import com.hhh.springai_test.common.ResultUtils;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/ai")
public class ChatController {

    @Resource
    private ZhiPuAiChatModel chatModel;

    // 存储每个用户的对话历史
    private final Map<String, List<String>> conversationHistory = new HashMap<>();

    @GetMapping("/generate")
    public BaseResponse<String> generate(@RequestParam(value = "msg", defaultValue = "Tell me a joke") String msg) {
        String message = this.chatModel.call(msg);
        return ResultUtils.success(message);
    }

    @GetMapping(value = "/generateStream",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message") String message,
                                             @RequestParam(value = "userId") String userId) {
        //获取初始化用户的对话历史
        List<String> userConversation = conversationHistory.computeIfAbsent(userId, k -> new ArrayList<>());
        //将用户输入的消息加入到对话历史中
        userConversation.add("User: " + message);
        String prompt = String.join("\n", userConversation) + "\nAI:";
        // 发起流式请求
        return chatModel.stream(new Prompt(new UserMessage(prompt)))
                .doOnTerminate(() -> {
                    // 在流结束时，将 AI 的回答添加到对话历史
                    String aiResponse = "AI: ";
                    userConversation.add(aiResponse);
                })
                .flatMap(Flux::just);
    }
}
